#pragma once

#include "compi/collectives.h"
#include "compi/context.h"
#include "compi/fixed_type.h"
#include "compi/result.h"
#include "compi/serializer.h"
#include "compi/tag.h"

#include <functional>
#include <memory>
#include <optional>
#include <string>
#include <type_traits>
#include <vector>

namespace compi
{

class TransportRequestImpl;

// 显式状态机：统一请求状态（需在请求类型前声明）
enum class RequestState
{
    Pending,   // 尚未进入传输（等待发布或配额）
    Inflight,  // 已发布至 MPI，进行中
    Completed, // 已完成（成功），可选择处理/释放
    Canceled,  // 已取消
    Failed     // 发生错误/异常
};

// 请求状态信息：用于 status() 返回当前或最近一次完成的元数据
struct RequestStatus
{
    int source = -1;                    // 源进程（未知为 -1）
    int tag = -1;                       // 消息标签值（未知为 -1）
    std::optional<MessageTag> tag_info; // 资源层标签信息（包含 ns_id）
    int count_bytes = 0;                // 负载字节数（若不可得则为 0）
    bool cancelled = false;             // 是否已被取消
    std::optional<std::string> error;   // 错误信息（若无错误则为空）
};

// 完成策略枚举（可按位组合）：用于控制完成后的动作顺序
enum CompletionStrategy : unsigned
{
    StrategyNone = 0u,
    StrategyUnpack = 1u << 0,  // 解包到用户 payload（针对序列化类型）
    StrategyProcess = 1u << 1, // 执行用户提供的处理回调
    /**
     * 释放内部资源（缓冲区、标签等）。
     *
     * 标签释放联动说明：
     * - 非阻塞原语（如 isend/irecv）在请求内部会对标签执行 retain；
     * - 完成时，若启用本策略位，将自动执行 release：
     *   优先使用命名空间句柄 `TagNamespace::release(tag)`；如无命名空间句柄，
     *   回退调用分配器 `MessageTagAllocator::release(tag)`；
     * - 如需禁用自动释放，请在调用 `set_strategy(...)` 时清除该位（示例：
     *   `req.set_strategy(req_strategy & ~StrategyRelease);`）。
     *
     * 完成语义与异常传播：
     * - 完成路径为 Unpack → Process → Release；
     * - 若 Process 阶段抛出异常，默认不执行 Release（以避免双重释放）；异常将沿 wait() 传播；
     * - 在不可传播异常的上下文（例如 poll() 的 noexcept），异常将被记录为错误消息，且默认不释放；
     * - 上层可在需要时调用 `TransportRequest::release()` 进行显式释放，该操作是幂等的。
     */
    StrategyRelease = 1u << 2
};

// 传输请求：非阻塞操作的占位与同步
class TransportRequest
{
public:
    // 构造一个空的传输请求。
    // 参数：无
    // 返回值：无
    // 异常：无
    TransportRequest() = default;

    // 使用底层实现指针构造传输请求。
    // 参数：impl 底层实现共享指针。
    // 返回值：无
    // 异常：无
    explicit TransportRequest(std::shared_ptr<TransportRequestImpl> impl) noexcept
        : _impl(std::move(impl))
    {
    }

    // 设置底层实现指针，用于通道层在发起非阻塞操作后回填请求对象。
    // 参数：impl 底层实现共享指针。
    // 返回值：无
    // 异常：无
    void set_impl(std::shared_ptr<TransportRequestImpl> impl) noexcept { _impl = std::move(impl); }

    // 非抛出等待：阻塞直到完成并返回统一的 Status。
    // 参数：无
    // 返回值：Status::OK 表示成功；错误时包含错误码与信息。
    // 异常：无（不抛出）。
    Result<void> wait() noexcept;

    // 查询是否已完成（非阻塞）。若尚未发布接收，将尝试使用非阻塞探测发布接收。
    // 参数：无
    // 返回值：true 表示已完成；false 表示未完成或出错。
    // 异常：无
    bool completed() noexcept;

    // 取消未完成的请求；若已完成或不存在底层实现则返回 false。
    // 参数：无
    // 返回值：true 表示取消成功并完成清理；false 表示无法取消（可能已完成或无请求）。
    bool cancel() noexcept;

    // 推进进度（非阻塞）：尝试匹配并发布接收、测试请求完成等；返回是否发生了可观察的进度。
    // 参数：无
    // 返回值：true 表示推进了进度（匹配/发布/完成）；false 表示无变化或出错。
    bool poll() noexcept;

    // 查询当前/最近一次完成的请求状态。
    // 参数：无
    // 返回值：RequestStatus，包含源、tag、字节数与错误/取消信息。
    RequestStatus status() const noexcept;

    // 查询当前/最近一次完成的请求状态，使用统一的Result类型。
    // 参数：无
    // 返回值：Result<RequestStatus>，成功时包含状态信息，失败时包含错误码和错误信息。
    // 异常：无
    Result<RequestStatus> result() const noexcept;

    // 设置完成后的处理回调（处理已解包的 payload 指针），用于“解包→处理→释放”中的处理阶段。
    // 参数：process 处理函数，接收 payload
    // 指针（固定类型为对象地址，序列化类型在解包后为对象地址）。
    void set_processor(std::function<void(void *)> process) noexcept;

    // 设置带上下文的处理回调：在完成时传入 payload 指针与请求状态元数据。
    // 参数：processor 处理函数，签名为 (void*, const RequestStatus&)
    // 说明：与 set_processor 共存，若两者均设置则按策略依次调用（先无上下文，再带上下文）。
    void set_processor(std::function<void(void *, const RequestStatus &)> processor) noexcept;

    // 添加一个完成后回调（在资源释放后调用）。
    // 参数：cb 无参回调。
    void add_callback(std::function<void()> cb);

    // 设置完成策略（按位组合）。默认值为 Unpack | Release。
    // 提示：清除 StrategyRelease 位可禁用自动标签释放，由调用方手动管理标签生命周期。
    void set_strategy(unsigned strategy_flags) noexcept;

    // 显式释放内部资源（缓冲区、标签等），与 StrategyRelease 行为一致且幂等。
    // 参数：无
    // 返回值：true 表示本次调用发生了实际释放；false 表示已释放或无可释放资源。
    // 异常：无
    bool release() noexcept;

    // 友元：批量等待管理器可访问内部实现以进行高效操作（如后续优化）。
    friend class RequestGroup;

private:
    std::shared_ptr<TransportRequestImpl> _impl;
};

// 批量请求组：提供 wait_all/test_some 等批量等待原语
class RequestGroup
{
public:
    // 添加一个请求到组
    void add(TransportRequest req) { _requests.emplace_back(std::move(req)); }

    // 当前组内请求数量
    std::size_t size() const noexcept { return _requests.size(); }

    // 阻塞等待所有请求完成；返回是否全部成功（失败策略可能抛出或终止）
    bool wait_all();

    // 非阻塞测试，返回当前已完成的请求索引集合；不会移除已完成请求。
    std::vector<std::size_t> test_some();

    // 推进所有请求的进度（匹配、发布、测试），返回是否有进度发生
    bool poll_all() noexcept;

private:
    std::vector<TransportRequest> _requests;
};

// 轻量工厂：为固定类型提供非阻塞发送/接收的请求构造
// 注意：这些工厂仅设置必要的字段以匹配 TransportRequest 的内部实现，
// 不涉及序列化类型。

// 固定类型：ANY_TAG 过滤接收（命名空间）
TransportRequest
irecv_fixed(Context & ctx, int src, void * payload, MPI_Datatype dt, TagNamespace ns) noexcept;

// 固定类型：命名空间标签下的非阻塞发送（可选同步模式）
TransportRequest isend_fixed(Context & ctx,
                             int dest,
                             const void * payload,
                             MPI_Datatype dt,
                             TagNamespace ns,
                             SendMode mode = SendMode::Default) noexcept;

// 固定类型：非阻塞发送（使用 TagLease 适配，避免双重释放）
TransportRequest isend_fixed(Context & ctx,
                             int dest,
                             const void * payload,
                             MPI_Datatype dt,
                             const TagLease & lease,
                             SendMode mode = SendMode::Default) noexcept;

// -----------------------------
// 序列化类型便捷工厂：字符串/字节向量
// -----------------------------
// ANY_TAG + 命名空间过滤的非固定长度接收（序列化负载）
TransportRequest
irecv_serialized(Context & ctx, int src, std::string & payload, TagNamespace ns) noexcept;

TransportRequest irecv_serialized(Context & ctx,
                                  int src,
                                  std::vector<std::byte> & payload,
                                  TagNamespace ns) noexcept;

// 命名空间标签下的非阻塞发送（序列化负载）
TransportRequest isend_serialized(Context & ctx,
                                  int dest,
                                  const std::string & payload,
                                  TagNamespace ns,
                                  SendMode mode = SendMode::Default) noexcept;

TransportRequest isend_serialized(Context & ctx,
                                  int dest,
                                  const std::vector<std::byte> & payload,
                                  TagNamespace ns,
                                  SendMode mode = SendMode::Default) noexcept;

// 序列化负载：非阻塞发送（使用 TagLease 适配，避免双重释放）
TransportRequest isend_serialized(Context & ctx,
                                  int dest,
                                  const std::string & payload,
                                  const TagLease & lease,
                                  SendMode mode = SendMode::Default) noexcept;

TransportRequest isend_serialized(Context & ctx,
                                  int dest,
                                  const std::vector<std::byte> & payload,
                                  const TagLease & lease,
                                  SendMode mode = SendMode::Default) noexcept;

// 模板便捷封装：对可序列化对象使用 serializer 将对象编码为字节并发送
template <typename T>
inline TransportRequest
isend_serialized(Context & ctx,
                 int dest,
                 const T & payload,
                 TagNamespace ns,
                 SendMode mode = SendMode::Default) noexcept
{
    auto writer = serialize_to_vector_writer(payload);
    // 复制到字节向量以与底层非模板工厂对接
    std::vector<std::byte> bytes(writer.data().begin(), writer.data().end());
    return isend_serialized(ctx, dest, bytes, ns, mode);
}

// 模板便捷封装：自动映射固定类型的 MPI_Datatype
template <typename T>
inline TransportRequest
irecv(Context & ctx, int src, T & payload, TagNamespace ns) noexcept
{
    if constexpr (types::is_fixed_type_v<T>)
    {
        return irecv_fixed(ctx, src, &payload, types::mpi_datatype_of<T>(), ns);
    }
    else if constexpr (std::is_same_v<T, std::string>)
    {
        return irecv_serialized(ctx, src, payload, ns);
    }
    else if constexpr (std::is_same_v<T, std::vector<std::byte>>)
    {
        return irecv_serialized(ctx, src, payload, ns);
    }
    else
    {
        // 通用序列化类型：先接收到字节缓冲，再在 Process 阶段反序列化到 T
        auto tmp_bytes = std::make_shared<std::vector<std::byte>>();
        TransportRequest req = irecv_serialized(ctx, src, *tmp_bytes, ns);
        req.set_processor([tmp_bytes, &payload](void *) { payload = deserialize<T>(*tmp_bytes); });
        req.set_strategy(StrategyUnpack | StrategyProcess | StrategyRelease);
        return req;
    }
}

template <typename T>
inline TransportRequest
isend(Context & ctx,
      int dest,
      const T & payload,
      TagNamespace ns,
      SendMode mode = SendMode::Default) noexcept
{
    if constexpr (types::is_fixed_type_v<T>)
    {
        return isend_fixed(ctx, dest, &payload, types::mpi_datatype_of<T>(), ns, mode);
    }
    else if constexpr (std::is_same_v<T, std::string>)
    {
        return isend_serialized(ctx, dest, payload, ns, mode);
    }
    else if constexpr (std::is_same_v<T, std::vector<std::byte>>)
    {
        return isend_serialized(ctx, dest, payload, ns, mode);
    }
    else
    {
        // 任意可序列化对象：使用 serializer 编码后发送
        return isend_serialized(ctx, dest, payload, ns, mode);
    }
}

} // namespace compi